package com.devyk.av.rtmp.library.stream

import android.media.MediaCodec
import com.devyk.av.rtmp.library.Contacts
import com.devyk.av.rtmp.library.utils.LogHelper
import java.nio.ByteBuffer
import java.util.ArrayList
import kotlin.experimental.and

/**
 * <pre>
 *     author  : devyk on 2020-07-17 15:43
 *     blog    : https://juejin.im/user/578259398ac2470061f3a3fb/posts
 *     github  : https://github.com/yangkun19921001
 *     mailbox : yang1001yk@gmail.com
 *     desc    : This is AnnexbHelper
 *     doc: 学习笔记: https://c2ixyjfcz7.feishu.cn/docx/GkrndT1t4oeh9vxWnf4cGQcXnnf?from=from_copylink
 * </pre>
 */
public class AnnexbHelper {

    // 定义一个监听器变量，用于监听NALU事件 (NALU 单独的帧)
    private var mListener: AnnexbNaluListener? = null

    // 存储PPS数据的字节数组
    private var mPps: ByteArray? = null

    // 存储SPS数据的字节数组
    private var mSps: ByteArray? = null

    // 标志是否需要上传PPS和SPS
    private var mUploadPpsSps = true

    /**
     * 用于存储Annexb搜索结果的内部类
     */
    internal inner class AnnexbSearch {
        // 起始码长度
        var startCode = 0

        // 是否匹配
        var match = false
    }

    // 定义AnnexbNaluListener接口
    interface AnnexbNaluListener {
        // 当检测到SPS和PPS时调用
        fun onSpsPps(sps: ByteArray?, pps: ByteArray?)

        // 当检测到视频帧时调用
        fun onVideo(data: ByteArray, isKeyFrame: Boolean)
    }

    // 设置AnnexbNaluListener监听器
    fun setAnnexbNaluListener(listener: AnnexbNaluListener) {
        mListener = listener
    }

    // 停止监听并清空缓存数据
    fun stop() {
        mListener = null
        mPps = null
        mSps = null
        mUploadPpsSps = true
    }

    /**
     * 处理硬编码生成的视频数据，将其解析为每一帧，并传递给FLV打包器
     * @param bb 硬编码后的数据缓冲区
     * @param bi 硬编码的BufferInfo
     */
    fun analyseVideoData(bb: ByteBuffer, bi: MediaCodec.BufferInfo) {
        // 设置缓冲区的起始和结束位置
        bb.position(bi.offset)
        bb.limit(bi.offset + bi.size)
        // 存储解析出来的帧
        val frames = ArrayList<ByteArray>()
        var isKeyFrame = false

        // 循环解析缓冲区中的数据
        while (bb.position() < bi.offset + bi.size) {
            // 解析出一个NALU
            val frame = annexbDemux(bb, bi)
            if (frame == null) {
                LogHelper.e(Contacts.TAG, "annexb not match.")
                break
            }

            // 如果是访问单元分隔符，跳过
            if (isAccessUnitDelimiter(frame)) {
                continue
            }

            // 处理PPS、SPS和普通视频帧
            when {
                isPps(frame) -> {
                    mPps = frame
                }

                isSps(frame) -> {
                    mSps = frame
                }

                else -> {
                    isKeyFrame = isKeyFrame(frame)
                    val naluHeader = buildNaluHeader(frame.size)
                    frames.add(naluHeader)
                    frames.add(frame)
                }
            }
        }
        // 如果有PPS和SPS，并且监听器不为空且需要上传PPS和SPS
        if (mPps != null && mSps != null && mListener != null && mUploadPpsSps) {
            if (mListener != null) {
                mListener!!.onSpsPps(mSps, mPps)
            }
            mUploadPpsSps = false
        }
        // 如果没有解析出帧或者监听器为空，返回
        if (frames.size == 0 || mListener == null) {
            return
        }
        // 计算所有帧的总大小
        var size = 0
        for (i in frames.indices) {
            val frame = frames[i]
            size += frame.size
        }
        // 创建一个字节数组存储所有帧数据
        val data = ByteArray(size)
        var currentSize = 0
        for (i in frames.indices) {
            val frame = frames[i]
            System.arraycopy(frame, 0, data, currentSize, frame.size)
            currentSize += frame.size
        }
        // 调用监听器的onVideo方法传递视频数据
        if (mListener != null) {
            mListener!!.onVideo(data, isKeyFrame)
        }
    }

    /**
     * 从硬编码的数据中提取一帧NALU
     * @param bb 硬编码后的数据缓冲区
     * @param bi 硬编码的BufferInfo
     * @return 提取到的NALU数据
     */
    private fun annexbDemux(bb: ByteBuffer, bi: MediaCodec.BufferInfo): ByteArray? {
        val annexbSearch = AnnexbSearch()
        avcStartWithAnnexb(annexbSearch, bb, bi)

        if (!annexbSearch.match || annexbSearch.startCode < 3) {
            return null
        }
        for (i in 0 until annexbSearch.startCode) {
            bb.get()
        }

        val frameBuffer = bb.slice()
        val pos = bb.position()
        while (bb.position() < bi.offset + bi.size) {
            avcStartWithAnnexb(annexbSearch, bb, bi)
            if (annexbSearch.match) {
                break
            }
            bb.get()
        }

        val size = bb.position() - pos
        val frameBytes = ByteArray(size)
        frameBuffer.get(frameBytes)
        return frameBytes
    }

    /**
     * 在硬编码的ByteBuffer中查找NALU
     * @param as 用于存储查找结果的AnnexbSearch对象
     * @param bb 硬编码后的数据缓冲区
     * @param bi 硬编码的BufferInfo
     */
    private fun avcStartWithAnnexb(`as`: AnnexbSearch, bb: ByteBuffer, bi: MediaCodec.BufferInfo) {
        `as`.match = false
        `as`.startCode = 0
        var pos = bb.position()
        while (pos < bi.offset + bi.size - 3) {
            // 如果前两个字节不是0x00，跳出循环
            if (bb.get(pos).toInt() != 0x00 || bb.get(pos + 1).toInt() != 0x00) {
                break
            }

            // 匹配到N[00] 00 00 01，其中N>=0
            if (bb.get(pos + 2).toInt() == 0x01) {
                `as`.match = true
                `as`.startCode = pos + 3 - bb.position()
                break
            }
            pos++
        }
    }

    // 构建NALU头部
    private fun buildNaluHeader(length: Int): ByteArray {
        val buffer = ByteBuffer.allocate(4)
        buffer.putInt(length)
        return buffer.array()
    }

    // 判断是否为SPS
    private fun isSps(frame: ByteArray): Boolean {
        if (frame.size < 1) {
            return false
        }
        // 5bits, 7.3.1 NAL unit syntax,
        // H.264-AVC-ISO_IEC_14496-10.pdf, page 44.
        //  7: SPS, 8: PPS, 5: I Frame, 1: P Frame
        val nal_unit_type = frame[0] and 0x1f
        return nal_unit_type.toInt() == SPS
    }

    // 判断是否为PPS
    private fun isPps(frame: ByteArray): Boolean {
        if (frame.size < 1) {
            return false
        }
        // 5bits, 7.3.1 NAL unit syntax,
        // H.264-AVC-ISO_IEC_14496-10.pdf, page 44.
        //  7: SPS, 8: PPS, 5: I Frame, 1: P Frame
        val nal_unit_type = frame[0] and 0x1f
        return nal_unit_type.toInt() == PPS
    }

    // 判断是否为关键帧
    private fun isKeyFrame(frame: ByteArray): Boolean {
        if (frame.size < 1) {
            return false
        }
        // 5bits, 7.3.1 NAL unit syntax,
        // H.264-AVC-ISO_IEC_14496-10.pdf, page 44.
        //  7: SPS, 8: PPS, 5: I Frame, 1: P Frame
        val nal_unit_type = frame[0] and 0x1f
        return nal_unit_type.toInt() == IDR
    }

    companion object {
        // 非IDR图像切片
        val NonIDR = 1

        // IDR图像切片
        val IDR = 5

        // 补充增强信息
        val SEI = 6

        // 序列参数集
        val SPS = 7

        // 图像参数集
        val PPS = 8

        // 访问单元分隔符
        val AccessUnitDelimiter = 9

        /**
         * 判断给定的帧是否为访问单元分隔符
         * @param frame 要检查的帧数据
         * @return 如果是访问单元分隔符，则返回true；否则返回false
         */
        private fun isAccessUnitDelimiter(frame: ByteArray): Boolean {
            // 如果帧的大小小于1，返回false
            if (frame.size < 1) {
                return false
            }
            // 5bits, 7.3.1 NAL单元语法,
            // 参考H.264-AVC-ISO_IEC_14496-10.pdf, 第44页。
            //  7: SPS, 8: PPS, 5: I Frame, 1: P Frame
            // 获取NAL单元类型
            val nal_unit_type = frame[0] and 0x1f
            // 如果NAL单元类型等于访问单元分隔符类型，返回true
            return nal_unit_type.toInt() == AccessUnitDelimiter
        }
    }
}

